import numpy as np
import torch
import faiss

from detectron2.data import MetadataCatalog

class GirRetrievalEvaluatorCenters():
    def __init__(self,centers_tensor):
        self._faiss_index = None
        self._recall_values = [1,5,10]
        self._features_dim = 4096
        self._centers = None
        self._center_faiss_index = None
        if centers_tensor is not None:
            centers_tensor = centers_tensor.cpu().numpy()
            self._centers = centers_tensor
            self._center_faiss_index = faiss.IndexFlatL2(centers_tensor.shape[-1])
            self._center_faiss_index.add(centers_tensor)
        pass
    
    def process_retrieval_ins_feat(self,instances,topk = 10):
        perd_feats = instances.pred_boxes_ins_feature.numpy()
        perd_feats_insflag = instances.pred_boxes_ins_mask.numpy()
        
        distances, predictions = self._center_faiss_index.search(perd_feats, max(self._recall_values))
        queries_num = len(perd_feats)
        for q in range(queries_num):
            # sort predictions by distance
            sort_idx = np.argsort(distances[q])
            predictions[q] = predictions[q, sort_idx]
            # remove duplicated predictions, i.e. keep only the closest ones
            
        pred_qids = predictions
        pred_qids[~perd_feats_insflag,:] = -1
        instances.pred_ins_ids = pred_qids
        pass
